package com.zyt.wiki.interceptor.rate;

import com.alibaba.fastjson.JSONObject;
import com.zyt.wiki.interceptor.rate.properties.RateTokenIpProperties;
import com.zyt.wiki.resp.CommonResp;
import com.zyt.wiki.util.RateUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 【令牌桶限流（Ip维度）】
 *  放在 非法参数过滤器 之后执行
 */
@Component
public class RateTokenByIpInterceptor implements HandlerInterceptor {

    public Logger logger = LoggerFactory.getLogger(RateTokenByIpInterceptor.class) ;

    // 缓存IP对应令牌桶数据（多线程安全）
    private ConcurrentHashMap<String, TokenBucket> cacheIpToken = new ConcurrentHashMap<String, TokenBucket>() ;

    @Autowired
    private RateTokenIpProperties properties; // Ip令牌桶配置

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (this.tryAcquire(request)) {
            // 令牌桶中有令牌，继续后续交易
            return true ;
        }else {
            this.logger.warn("@@@@@@@@ Access denied! rate-ip limiter: {} {}", request.getMethod(), request.getRequestURI()) ;
            // 返回特定报文给前端
//            response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value()) ;// 为方便前端VUE获取报文信息进行弹框提示，这里就先返回200，不返回429了
            response.setStatus(HttpStatus.OK.value()) ;
            response.setContentType("application/json;charset=UTF-8");
            response.setCharacterEncoding("UTF-8");
            CommonResp<String> commonResp = new CommonResp<>() ;
            commonResp.setSuccess(false) ;
            commonResp.setMessage("请求太频繁，请稍后再试");
            response.getWriter().print(JSONObject.toJSON(commonResp));
            return false ; // 返回false，则请求中断，不再往后到 其他拦截器、Controller
        }
    }

    /** true：能通过    false：不能通过 */
    public boolean tryAcquire(HttpServletRequest request){
        // 开关关闭时，直接通过
        if(!properties.isEnabled()){
            return true ;
        }
        // 若客户端IP为空，则直接通过
        String ip = RateUtils.getIP(request) ;
        if(StringUtils.isEmpty(ip)){
            return true ;
        }
        // 从缓存中获取令牌数据
        TokenBucket tokenBucket = cacheIpToken.get(ip) ;
        if(null == tokenBucket){
            // 第一次请求，默认初始化令牌桶（要进行同步控制）
            synchronized (this) {
                TokenBucket tokenBuc =cacheIpToken.get(ip) ;
                if(null == tokenBuc){    // 并发场景下，需二次判断，否则排队后的线程，后续仍会进来创建新对象
                    tokenBuc = new TokenBucket("token-ip", properties.getBucketLimit(), properties.getTokenPerInterval(), properties.getIntervalInMills()) ;
                    cacheIpToken.put(ip, tokenBuc) ;
                }
                return tokenBuc.tryAcquire() ;
            }
        }else {
            // 非第一次请求
            return tokenBucket.tryAcquire() ;   // 该方法已做同步控制
            // 更新tokenBucket到缓存中 【上面用的引用变量，会直接影响原缓存Map中的值，无需再put一遍】
        }
    }

}
